#ifndef EST_POLYA_SGD_HEADER
#define EST_POLYA_SGD_HEADER

#include<string>
#include<fstream>
#include<sstream>
#include<iostream>
#include<algorithm>
#include<iterator>
#include<vector>
#include<map>
#include<cstdlib>
#include<ctime>
#include<cmath>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_sort.h>
#include <gsl/gsl_statistics.h>
#include <gsl/gsl_sf_psi.h>
#include "def.hpp"
#include "utils.hpp"
#include "io.hpp"
#include "dist_polya.hpp"

using namespace std;

enum Polya_SGD_Type{
  Polya_SGD_SUCCESS,
  Polya_SGD_BAD
};

typedef struct Polya_SGD_Search_Struct{
  Polya_SGD_Type search_type;
  double search_value;
}Polya_SGD_Search_Type;

Polya_SGD_Type est_polya_sgd_iter(int**& d, int*& n, double*& param, int DIM_N, int DATA_N, int*& buffer_id, double rate){  
  for(int i=0;i < DATA_N; i++){
    double ParamSum = get_sum(param,DIM_N);
    for(int k=0; k < DIM_N; k++){
      double g = gsl_sf_psi(ParamSum) - gsl_sf_psi(n[buffer_id[i]] + ParamSum) + gsl_sf_psi(d[buffer_id[i]][k] + param[k]) - gsl_sf_psi(param[k]);
      double a =  param[k] + rate * g;
      if(a > 0){
	param[k] = a;
      }
	else{
	  return Polya_SGD_BAD;
	}
    }
  }
  return Polya_SGD_SUCCESS;
}

Polya_SGD_Search_Type est_polya_sgd_search_gain(int**& d, int*& n, double*& param, int DIM_N, int DATA_N, int*& buffer_id, double base_rate, int iter){
  double j[] = {0.1,0.111111,0.125,0.14285,0.167,0.2,0.25,0.333,0.5,0,1,2,3,4,5,6,7,8,9,10};
  double* backup_param = new double[DIM_N];
  double* my_param = new double[DIM_N];
  double largest_gain = 0.0;
  Polya_SGD_Search_Type result;
  int failed = 0;
  
  result.search_value = -1.0;
  result.search_type = Polya_SGD_BAD;

  for(int k=0;k<DIM_N;k++){
    backup_param[k] = param[k];
  }

  double current_likelihood = polya_likelihood(d, n, backup_param, DIM_N, DATA_N ,buffer_id);  
  for(int i=0; i < 20; i ++){
    for(int k=0; k < DIM_N; k ++){
      my_param[k] = backup_param[k];
    }
    double rate = base_rate / (j[i] + iter);
    Polya_SGD_Type current_type = est_polya_sgd_iter(d, n, my_param, DIM_N, DATA_N, buffer_id, rate);
    if(current_type == Polya_SGD_BAD){
      failed ++;
    }
    else{
      double after_likelihood = polya_likelihood(d, n, my_param, DIM_N, DATA_N, buffer_id);
      double current_gain = after_likelihood - current_likelihood;
      //      cout << "\t trying:" << j[i] << " " << current_gain << endl;
      if(current_gain > largest_gain){
	largest_gain = current_gain;
	result.search_value = j[i];
	result.search_type = Polya_SGD_SUCCESS;
      }
    }
  }
  delete[] backup_param;
  delete[] my_param;
  //  cout << "\t chosen:" << result.search_value << endl;
  return result;
}

void est_polya_sgd(gsl_rng* r, int**& d, int*& n, double*& param, int DIM_N, int DATA_N){
  int SAMPLE_NUM = 1000;
  int* all_id = new int[DATA_N];
  int* buffer_id = new int[SAMPLE_NUM];
  for(int i=0; i < DATA_N; i++){
    all_id[i] = i;
  }

  double rate = PuzaDEF::Instance()->EST_Polya_SGD_Initial_Rate;
  double new_rate = rate;

  int iter  = 0;
  double before_norm = 0.0;
  double after_norm = 0.0;
  
  while(1){
    //    cout << "Current Rate:" << rate << endl;
    gsl_ran_shuffle (r, all_id, DATA_N, sizeof (int));
    gsl_ran_choose(r, buffer_id, 1000, all_id, DATA_N, sizeof(int));
    Polya_SGD_Search_Type search_result = est_polya_sgd_search_gain(d, n, param, DIM_N, SAMPLE_NUM, buffer_id, rate, iter);
    if(search_result.search_type == Polya_SGD_SUCCESS){
      new_rate = rate / (search_result.search_value + iter);
      before_norm = get_sum(param,DIM_N);
      Polya_SGD_Type iter_result = est_polya_sgd_iter(d, n, param, DIM_N, DATA_N, all_id, new_rate);
      if(iter_result == Polya_SGD_SUCCESS){
	after_norm = get_sum(param,DIM_N);
	double error = fabs(after_norm - before_norm);
	if(error < 1e-8){
	  //  cout <<"MEET EXIT CONDITION" << endl;                                                  
	  break;
	}
	iter ++;
	//	cout << "ITERATION :" << iter << " Objective:" << polya_likelihood(d, n, param, DIM_N, DATA_N, all_id) << endl;
	rate = new_rate;
	if(iter > 20){
	  break;
	}
      }
      else{
	// adjust rate     
	rate = rate / 10;
	//	cout << "Adjust Rate:" << rate << endl;
      }
    }
    else{
      // adjust rate
      rate = rate / 10;
      //      cout << "Adjust Rate:" << rate << endl;
    }
    
  }
  delete[] all_id;
  delete[] buffer_id;
}

#endif
